# Copyright (c) 2014 City of Hope
# Authors: Juan Du and Dustin E Schones
#!/usr/bin/python

""" Outline:
1. read in FAIRE reads file (BED3/BED6/BEDPE), store read coordinates in dictionary 

2. getScoreInWindow: scan genome with 250bp window (50bp increments), count reads in each window,
   store counts in dictionary (all_scores) define threshold for peak cut-off. Input from user: p-value.

3. getEnrichedRegion: scan genome to find windows with reads counts more than threshold. Continuous
   windows merge into one enriched region.

   dynamic background (the neighboring peak within background will be normalized to the global threshold)

   Multi-scle background model: using 10k, 20k and 50k BG, and choose the one bigger than,
   but closest to chromosome-wide threshold.
"""

from __future__ import division
import re, os, sys, shutil
from math import *
from string import *
from optparse import OptionParser
import operator
from bisect import *
import numpy as np
from scipy.stats import nbinom
from scipy.special import psi
from scipy.optimize import brentq

import BED;
import GenomeData;


class Coords:
    def __init__(self, start, end):
        self.start = start
        self.end = end
    def __set__(self, start, end):
        self.start = start
        self.end = end
    def getCoord(self):
        outstring=str(self.start) + "\t" + str(self.end)
        try:
            return outstring
        except:
            sys.stderr.write("No coord information for %s\n" % self)
            return ''


def getScore(tags, tag_index, pos_start, pos_end):
    n = 0;
    for i in tag_index:
        if tags[i].start <= pos_start and tags[i].end <= pos_start: n = n;
        elif tags[i].start >= pos_end and tags[i].end >= pos_end: n = n;
        else: n += 1;
    return n;

def getTagPos(tags):
    pos = [];
    for t in tags:
        pos.append(t.start)
    return pos;

def getTagsInIndex(tag_pos, pos_start, pos_end):
    i_up = pos_start - 1000;
    i_down = pos_end + 1000;
    start_ind = bisect_left(tag_pos, i_up);
    end_ind = bisect_right(tag_pos, i_down);
    return range(start_ind, end_ind);

def getScoreInWindow(window_size, step, all_tags, species):
    all_scores = {};
    for chrom in all_tags:
        if len(all_tags[chrom]) != 0:
            all_scores[chrom]=[];
            pairs = all_tags[chrom];
            if len(pairs) > 0:
                pairs.sort(key=operator.attrgetter('start'));
                tag_pos = getTagPos(pairs);
                pos_start = 0;
                end = GenomeData.species_chrom_lengths[species][chrom];
                n_steps = int(float(end)-(window_size-step))/float(step);
                pos_end = pos_start + window_size -1;
                tag_index = getTagsInIndex(tag_pos, pos_start, pos_end);
                score = getScore(pairs, tag_index, pos_start, pos_end)
                all_scores[chrom].append(score);
                for i in range(int(n_steps) - 1):
                    pos_start += step;
                    pos_end += step;
                    tag_index = getTagsInIndex(tag_pos, pos_start, pos_end);
                    score = getScore(pairs, tag_index, pos_start, pos_end)
                    all_scores[chrom].append(score);
    return all_scores;


def getEnrichedRegionDynamicBG(all_scores, threshold, window_size, step, pvalue, species, bgsizes):
    regions = BED.BED(species);
    for chrom in all_scores:
        scores = all_scores[chrom];
        j = 0;

        (bg1, bg2, bg3) = bgsizes.split(",");
        background1 = atoi(bg1)
        background2 = atoi(bg2)
        background3 = atoi(bg3)

        chrom_end = GenomeData.species_chrom_lengths[species][chrom];
        for i in range(len(scores)):
            if int(scores[i]) >= threshold and j == 0:
                can_reg_start = i;
                j += 1;
            elif int(scores[i]) < threshold and j != 0:
                can_reg_end = i;
                can_reg_mid = (can_reg_start + can_reg_end)/2;

                bg1_scores = scores[ int(max(0,(can_reg_mid - (background1/2)/step))) : int(min((can_reg_mid + (background1/2)/step),chrom_end)) ];
                n_bg1_scores = normalizeBG(bg1_scores, threshold);
                bg1_threshold = getThreshold(pvalue, n_bg1_scores);

                bg2_scores = scores[ int(max(0,(can_reg_mid - (background2/2)/step))) : int(min((can_reg_mid + (background2/2)/step),chrom_end)) ];
                n_bg2_scores = normalizeBG(bg2_scores, threshold);
                bg2_threshold = getThreshold(pvalue, n_bg2_scores);

                bg3_scores = scores[ int(max(0,(can_reg_mid - (background3/2)/step))) : int(min((can_reg_mid + (background3/2)/step),chrom_end)) ];
                n_bg3_scores = normalizeBG(bg3_scores, threshold);
                bg3_threshold = getThreshold(pvalue, n_bg3_scores);

                opt_threshold = getOptimizedThreshold(threshold, bg1_threshold, bg2_threshold, bg3_threshold);

                if np.mean(scores[ int(can_reg_start) : int(can_reg_end)]) >= opt_threshold:
                    peak_start = int( can_reg_start*step + window_size/2 + 1 );
                    peak_end = int( can_reg_end*step + window_size/2 );
                    if int( peak_end - peak_start ) >= 100:
                        coord = BED.BED3(chrom, peak_start, peak_end);
                        regions[chrom].append(coord);
                j = 0;
    return regions;


def getOptimizedThreshold(threshold, bg1_thr, bg2_thr, bg3_thr):
    if threshold >= max(bg1_thr, bg2_thr, bg3_thr):
        opt_threshold = threshold;
    elif bg1_thr > threshold and (bg2_thr >= bg1_thr or bg2_thr <= threshold) and (bg3_thr >= bg1_thr or bg3_thr <= threshold):
        opt_threshold = bg1_thr;
    elif bg2_thr > threshold and (bg1_thr >= bg2_thr or bg1_thr <= threshold) and (bg3_thr >= bg2_thr or bg3_thr <= threshold):
        opt_threshold = bg2_thr;
    elif bg3_thr > threshold and (bg1_thr >= bg3_thr or bg1_thr <= threshold) and (bg2_thr >= bg3_thr or bg2_thr <= threshold):
        opt_threshold = bg3_thr;
    return opt_threshold;


def normalizeBG(bg_scores, threshold):
    n_bg_scores = [];
    for i in range(len(bg_scores)):
        if bg_scores[i] <= threshold:
            n_bg_scores.append(bg_scores[i]);
        else:
            n_bg_scores.append(int(threshold));
    return n_bg_scores;


def getThreshold(p_value, scores):
    try:
        (size, prop) = fit_nbinom(np.array(scores));
        threshold = nbinom.ppf(1-p_value, size, prop);
    except:
        threshold = 0;
    return threshold;

# the fit_nbinom funtion is from: https://gist.github.com/gjx/5987413
def fit_nbinom(k):
    N = len(k)
    n = brentq(lambda r: sum(psi(k + r)) - N*psi(r) + N*log(r/(r + sum(k/N))),
               np.finfo(np.float).eps,
               np.max(k))
    p = n/(n + sum(k/N))
    return n, p



def printOut(coords, outfilename):
    outfile = open(outfilename, 'w');
    for chrom in coords.keys():
        for p in coords[chrom]:
            outline = p.chrom + "\t" + str(p.start) + "\t" + str(p.end) + "\n";
            outfile.write(outline);


def main(argv):
    parser = OptionParser()
    parser.add_option("-i", "--inputfile", action="store", type="string",
                      dest="infile", help="input file with tag coords", metavar="<file>")
    parser.add_option("-f", "--inputfileformat", action="store", type="string",
                      dest="infileformat", help="input file format (BED3/BED6/BEDPE), default: BED3", metavar="<str>", default="BED3")
    parser.add_option("-o", "--outfile", action="store", type="string",
                      dest="outfile", help="the name of the output BED file", metavar="<file>")
    parser.add_option("-s", "--species", action="store", type="string",
                      dest="species", help="species, default: hg19", metavar="<str>", default="hg19")
    parser.add_option("-p", "--p-value", action="store", type="float",
                      dest="pvalue", help="p-value for searching potential enriched window, default: 0.02", metavar="<float>", default=0.02)
    parser.add_option("-g", "--gthreshold", action="store", type="float",
                      dest="gthreshold", help="The genomewide threshold calculated from sampled windows (required)", metavar="<float>")
    parser.add_option("-w", "--windowsize", action="store", type="int",
                      dest="window", help="window size for scan genome, default: 250", metavar="<int>", default=250)
    parser.add_option("-t", "--step", action="store", type="int",
                      dest="step", help="window increment, default: 50", metavar="<int>", default=50)
    parser.add_option("-b", "--bgsizes", action="store", type="string",
                      dest="bgsizes", help="BG sizes, comma delimited, default: 10000,20000,50000", metavar="<str>", default="10000,20000,50000")

    (opt, args) = parser.parse_args(argv)
    if len(argv) < 4:
        parser.print_help()
        sys.exit(1)

    if opt.infileformat == "BEDPE":
        tag_coords = BED.BEDPE(opt.species, opt.infile, "BPAIR");
    else:
        tag_coords = BED.BED(opt.species, opt.infile, opt.infileformat);

    all_scores = getScoreInWindow(opt.window, opt.step, tag_coords, opt.species);
    regions = getEnrichedRegionDynamicBG(all_scores, opt.gthreshold, opt.window, opt.step, opt.pvalue, opt.species, opt.bgsizes);
    printOut(regions, opt.outfile);


if __name__ == "__main__":
    main(sys.argv)
